from scipy.spatial.transform import Rotation as R
from utils import path_util
import joblib
import json
import numpy as np
import metric


def eval_image(pred_poses, idx):
    lidar_to_camera_RT = np.array([-0.0043368991524, -0.99998911867, -0.0017186757713, 0.016471385748, -0.0052925495236, 0.0017416212982, -
                                   0.99998447772, 0.080050847871, 0.99997658984, -0.0043277356572, -0.0053000451695, -0.049279053295, 0, 0, 0, 1]).reshape(4, 4)
    lidar_to_camera_R = lidar_to_camera_RT[:3, :3]
    camera_to_lidar_R = np.linalg.inv(lidar_to_camera_R)

    for pred_pose in pred_poses:
        pred_pose[:3] = (R.from_matrix(camera_to_lidar_R) *
                         R.from_rotvec(pred_pose[:3])).as_rotvec()

    pose_folder = '/home/ljl/hdd/lidarcap/labels/3d/pose/{}'.format(idx)
    gt_poses = []
    for pose_filename in filter(lambda x: x.endswith('.json'), path_util.get_sorted_filenames_by_index(pose_folder)):
        with open(pose_filename) as f:
            content = json.load(f)
            gt_pose = np.array(content['pose'], dtype=np.float32)
            gt_poses.append(gt_pose)
    gt_poses = np.stack(gt_poses)

    assert len(pred_poses) == len(gt_poses)

    metric.output_metric(pred_poses, gt_poses)


def eval_vibe():
    for idx in [7, 24, 29, 41]:
        print(idx)
        vibe_result_file = '/home/ljl/tmp/vibe/{}_vibe_output.pkl'.format(idx)
        vibe_result = joblib.load(vibe_result_file)

        assert len(vibe_result) == 1

        vibe_result = vibe_result[list(vibe_result.keys())[0]]
        pred_poses = vibe_result['pose']

        eval_image(pred_poses, idx)


def eval_hmr():
    for idx in [7, 24, 29, 41]:
        print(idx)
        hmr_result_file = '/home/ljl/hmr/output/{}_hmr.npy'.format(idx)
        hmr_result = np.load(hmr_result_file)
        if idx == 24 or idx == 29:
            hmr_result = hmr_result[:-1]
        hmr_result = hmr_result.squeeze()
        eval_image(hmr_result, idx)


if __name__ == '__main__':
    eval_hmr()
